package pool

import (
	"context"
	"errors"
	"sync"
	"yunpan/store/ceph"
)

// BaseWorkerPool 基础工作池流程
type BaseWorkerPool struct {
	workers   int //携程数
	taskQueue chan Task
	ctx       context.Context //控制工作池的生命周期
	cancel    context.CancelFunc
	wg        sync.WaitGroup //等待工作协程完成
	mu        sync.RWMutex   //保护工作池状态
	isRunning bool           //运行状态判断
}

// 创建新的工作池
func NewWorkerPool(workers int, quueSize int) *BaseWorkerPool {
	ctx, cancel := context.WithCancel(context.Background())
	return &BaseWorkerPool{
		workers:   workers,
		taskQueue: make(chan Task, quueSize),
		ctx:       ctx,
		cancel:    cancel,
		isRunning: false,
	}
}

// 启动工作池
func (p *BaseWorkerPool) Start() {
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.isRunning {
		return
	}

	p.isRunning = true

	//启动工作协程
	for i := 0; i < p.workers; i++ {
		p.wg.Add(1)
		go p.worker()
	}
}

// 停止工作池
func (p *BaseWorkerPool) Stop() {
	p.mu.Lock()
	if !p.isRunning {
		p.mu.Unlock()
		return
	}
	p.isRunning = false
	p.mu.Unlock()

	//取消上下文。关闭任务队列，等待所有工作协程完成
	p.cancel()
	close(p.taskQueue)
	p.wg.Wait()
}

// 提交任务
func (p *BaseWorkerPool) Submit(task Task) error {
	p.mu.RLock()
	if !p.isRunning {
		p.mu.RUnlock()
		return errors.New("worker pool is not running")
	}
	p.mu.RUnlock()

	select{
	case p.taskQueue <- task:
		return nil
	case <- p.ctx.Done():
		return p.ctx.Err()
	}
}

//工作协程
func (p *BaseWorkerPool) worker(){
	defer p.wg.Done()
	for{
		select{
		case <-p.ctx.Done():
			return
		case task,ok := <-p.taskQueue:
			if !ok{
				//任务队列关闭
				return
			}
			if err := task.Execute();err != nil{
				//任务失败处理
				task.OnFailure(err)
				//如果需要重试且未超过最大重试次数，重新提交任务
				if task.GetRetryCount() > 0{
					p.Submit(task)
				}
			}	
		}
	}
}

//ceph专用线程池
type CephWorkerPool struct{
	*BaseWorkerPool
	connPool sync.Pool //ceph连接池
}

//创建ceph工作池
func NewCephWorkerPool(workers int,queueSize int) *CephWorkerPool{
	pool := &CephWorkerPool{
		BaseWorkerPool: NewWorkerPool(workers,queueSize),
	}

	//初始化连接池
	pool.connPool.New = func() interface{}{
		return &ceph.StorageSession{}
	}

	return pool
}

//数据库专用工作池
type DBWorkerPool struct {
    *BaseWorkerPool
}

// NewDBWorkerPool 创建数据库工作池
func NewDBWorkerPool(workers int, queueSize int) *DBWorkerPool {
    return &DBWorkerPool{
        BaseWorkerPool: NewWorkerPool(workers, queueSize),
    }
}

// HashWorkerPool 哈希计算专用工作池
type HashWorkerPool struct {
    *BaseWorkerPool
}

// NewHashWorkerPool 创建哈希计算工作池
func NewHashWorkerPool(workers int, queueSize int) *HashWorkerPool {
    return &HashWorkerPool{
        BaseWorkerPool: NewWorkerPool(workers, queueSize),
    }
}